{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "21743397",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "fd9c493c",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.ones(16, 3, 49, 32)     # B, num_heads, N, C//num_heads)\n",
    "y = torch.randn(16, 3, 49, 32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b1c320e4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[1., 1., 1.,  ..., 1., 1., 1.],\n",
      "        [1., 1., 1.,  ..., 1., 1., 1.],\n",
      "        [1., 1., 1.,  ..., 1., 1., 1.],\n",
      "        ...,\n",
      "        [1., 1., 1.,  ..., 1., 1., 1.],\n",
      "        [1., 1., 1.,  ..., 1., 1., 1.],\n",
      "        [1., 1., 1.,  ..., 1., 1., 1.]])\n",
      "tensor([[-8.4423e-01,  1.5935e-01,  1.2054e+00,  ...,  1.3424e+00,\n",
      "         -2.9673e-01,  2.4951e-01],\n",
      "        [ 1.4211e-01, -7.4567e-01, -9.7593e-01,  ..., -1.0812e-02,\n",
      "          1.1873e+00, -8.3383e-02],\n",
      "        [ 7.3454e-02, -2.3788e+00,  2.7886e+00,  ...,  4.6028e-01,\n",
      "         -1.7123e+00, -2.1806e-01],\n",
      "        ...,\n",
      "        [ 4.1885e-02,  5.1768e-01,  1.0281e-01,  ...,  1.0447e+00,\n",
      "          6.4074e-01,  1.5139e-02],\n",
      "        [-1.3809e+00,  5.2040e-01, -2.3739e+00,  ..., -1.9111e+00,\n",
      "         -1.5826e+00, -2.5849e-01],\n",
      "        [-1.4000e+00, -9.1961e-01,  1.1305e+00,  ...,  2.4868e-04,\n",
      "         -1.2108e+00,  1.4915e-01]])\n"
     ]
    }
   ],
   "source": [
    "print(x[0][0])\n",
    "print(y[0][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f3b8f202",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([16, 3, 49, 49])\n",
      "tensor([[ -1.3152,  -1.1164,  -5.7600,  ...,   2.1797, -18.1372,  -6.2213],\n",
      "        [ -1.3152,  -1.1164,  -5.7600,  ...,   2.1797, -18.1372,  -6.2213],\n",
      "        [ -1.3152,  -1.1164,  -5.7600,  ...,   2.1797, -18.1372,  -6.2213],\n",
      "        ...,\n",
      "        [ -1.3152,  -1.1164,  -5.7600,  ...,   2.1797, -18.1372,  -6.2213],\n",
      "        [ -1.3152,  -1.1164,  -5.7600,  ...,   2.1797, -18.1372,  -6.2213],\n",
      "        [ -1.3152,  -1.1164,  -5.7600,  ...,   2.1797, -18.1372,  -6.2213]])\n"
     ]
    }
   ],
   "source": [
    "z = (x @ y.transpose(-2, -1))\n",
    "print(z.size())\n",
    "print(z[0][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "fe97a060",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([16, 3, 49])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test = torch.sum(y, dim=-1)\n",
    "test.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "044e7a60",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ -1.3152,  -1.1164,  -5.7600,   2.4118,  12.0723,   4.6559,  -0.3333,\n",
       "          4.7077,  -3.3706,   0.0838,  -0.2500,  -1.8586,  -4.8678,  -5.7657,\n",
       "         -3.4342,  -3.9897,  -6.5426,   2.8364,  -0.2459,   3.3895,  -1.0704,\n",
       "         -3.2284,  -1.2611,  -5.4507,   0.4418,  -7.8946,   6.3007,   6.0482,\n",
       "          2.3416,  -3.7291,  -7.7350,   0.3407,  -2.2498,  10.5512,  10.0421,\n",
       "         -8.7029,  -0.4234,  -9.8404,  -4.8469,  -7.0324,   7.0553,   2.1268,\n",
       "         -0.1929,   1.0378,  -3.7582,  -1.4902,   2.1797, -18.1372,  -6.2213])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test[0][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "502f209f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([16, 3, 32, 49])\n",
      "torch.Size([16, 3, 49, 32])\n"
     ]
    }
   ],
   "source": [
    "x_transpose = x.transpose(-2, -1)\n",
    "print(x_transpose.shape)\n",
    "print(x.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "1c9cbe2a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3.,\n",
      "        3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3.,\n",
      "        3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3.])\n"
     ]
    }
   ],
   "source": [
    "x_transpose += 1.\n",
    "print(x_transpose[0][0][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "187ad8ba",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3.,\n",
      "        3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3.])\n"
     ]
    }
   ],
   "source": [
    "print(x[0][0][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76d3bed1",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
